package com.mfs.proxy.core.client;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.mfs.common.utills.ChannelUtil;
import com.mfs.proxy.core.client.config.ClientConfig;
import com.mfs.proxy.core.enums.MessageTypeEnum;
import com.mfs.proxy.core.util.ProxyTaskExecutor;
import com.mfs.proxy.core.message.mailbox.CommunicationMailbox;
import com.mfs.proxy.core.server.config.ServerConfig;
import com.mfs.proxy.core.server.endpoint.ProxyEndpoint;
import lombok.Data;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * 代理客户端对象，代理服务端持有
 * 可以通过改对象实现与客户端的通信、代理任务的执行
 * @author mfs
 * @date 2023/7/1 14:23
 **/
@Data
public class ProxyClient implements Closeable {
    private final static Logger logger = LoggerFactory.getLogger(ProxyClient.class);


    /**
     * 客户端配置
     * */
    private ClientConfig clientConfig;

    /**
     * 客户端通讯邮箱
     * */
    private CommunicationMailbox communicationMailbox;

    /**
     * 代理端点
     * */
    private ProxyEndpoint proxyEndpoint;

    /**
     * 代理任务执行器
     * */
    private ProxyTaskExecutor executor;

    /**
     * 正在处理的请求，防止userRequestCache移除过期通道时错误关闭了正在使用的通道
     * */
    private Set<String> processingRequest = Collections.synchronizedSet(new HashSet<>());

    /**
     * 用户请求缓存
     * key: requestNo, 具备超时机制
     * */
    private Cache<String, SocketChannel> userRequestCache  = Caffeine.newBuilder()
            .initialCapacity(128)
            .maximumSize(1024)
            .expireAfterWrite(ServerConfig.getInstance().getRequestTimeout(), TimeUnit.MINUTES)
            .evictionListener(((key, value, cause) -> {
                if (!processingRequest.contains(key)) {
                    logger.info("请求：{}超时将被关闭", key.toString());
                    ChannelUtil.closeSocketChannel((SocketChannel) value);
                }
            }))
            .build();

    public ProxyClient(ClientConfig clientConfig, CommunicationMailbox communicationMailbox, ProxyEndpoint proxyEndpoint, ProxyTaskExecutor executor) {
        this.clientConfig = clientConfig;
        this.communicationMailbox = communicationMailbox;
        this.proxyEndpoint = proxyEndpoint;
        this.executor = executor;
    }

    /**
     * 收到用户请求，先缓存到对应的代理客户端对象的缓存中
     * 然后向代理客户端发送一条通道请求消息（含一个随机的请求编号）
     *
     * @param userSocketChannel 用户请求通道
     * */
    public void processRequest(SocketChannel userSocketChannel) {
        String requestNo = UUID.randomUUID().toString();
        logger.info("准备处理代理请求：{}......", requestNo);
        userRequestCache.put(requestNo, userSocketChannel);

        try {
            communicationMailbox.send(MessageTypeEnum.CHANNEL_REQUEST, requestNo);
        } catch (IOException e) {
            logger.error("发送请求通道消息失败", e);
            try {
                userSocketChannel.close();
            } catch (IOException ioException) {
                logger.error("用户通道已关闭", e);
            }
        }
    }

    /**
     * 收到代理客户端的通道请求响应（使用新建立的代理通带来完成的，并不是使用代理客户端与代理服务端的邮箱通道）
     * 通道响应中会包含一个请求编号，该编号由通道请求消息发送带代理客户端，代理客户端响应时回传回来。
     *
     * 使用改请求编号获取缓存中的用户请求通道，建立用户请求通道、代理通道之间的一一对应关系。
     * 创建一个新的线程，用于处理用户请求
     *
     * @param requestNo 请求编号
     * @param proxyChannel 代理通道
     * */
    public void brokered(String requestNo, SocketChannel proxyChannel) {
        logger.info("正在处理代理请求：{}......", requestNo);
        SocketChannel userChannel = userRequestCache.getIfPresent(requestNo);
        if (userChannel == null) {
            ChannelUtil.closeSocketChannel(proxyChannel);
            logger.error("未找到对应的用户通道，可能改用户请求已经超时，直接关闭代理通道");
            return;
        }
        userRequestCache.invalidate(requestNo);
        processingRequest.add(requestNo);
        // 创建新线程处理用户请求
        executor.execute(() -> {
            Selector selector = null;
            try {
                try {
                    selector = Selector.open();
                } catch (IOException e) {
                    logger.error("打开筛选器失败，用户通道、代理通道将被关闭，不再处理改请求：{}",requestNo, e);
                    return;
                }

                SelectionKey proxyKey = null;
                SelectionKey userKey = null;
                try {
                    proxyKey = proxyChannel.register(selector, SelectionKey.OP_READ);
                    userKey = userChannel.register(selector, SelectionKey.OP_READ);
                    while (!Thread.currentThread().isInterrupted()) {
                        selector.select();
                        Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
                        while (iterator.hasNext()) {
                            SelectionKey next = iterator.next();
                            byte[] content = ChannelUtil.read((SocketChannel)next.channel());
                            if (content.length < 1) {
                                userKey.cancel();
                                proxyKey.cancel();
                                return;
                            }
                            if (next == proxyKey) {
                                userChannel.write(ByteBuffer.wrap(content));
                            } else if (next == userKey) {
                                proxyChannel.write(ByteBuffer.wrap(content));
                            }
                            iterator.remove();
                        }
                    }
                } catch (Exception e) {
                    logger.error("代理请求处理失败, 用户通道、代理通道将被关闭，不再处理改请求：{}", requestNo, e);
                    return;
                }
            } finally {
                logger.info("代理：{}请求处理完成", requestNo);
                ChannelUtil.closeSocketChannel(proxyChannel, userChannel);
                if (selector != null) {
                    try {
                        selector.close();
                    } catch (IOException e) {
                    }
                }
                processingRequest.remove(requestNo);
            }

        });
    }

    @Override
    public void close() throws IOException {
        communicationMailbox.close();
        proxyEndpoint.close();
    }

    @Override
    protected void finalize() throws Throwable {
        try {
            logger.debug("客户端销毁中");
            close();
        } catch (IOException e) {

        } finally {
            super.finalize();
        }
    }
}
